package algorithm.shipinke;

import java.util.Random;

/**
 * 蓄水池采样
 * 给定一个数据流，数据流长度N很大，且N直到处理完所有数据之前都不可知，请问如何在只遍历一遍数据（O(N)）的情况下，能够随机（每个数据被选到的概率等可能）选取出m个不重复的数据。
 *
 * 重点：
 * 数据流长度N很大且不可知，所以不能一次性存入内存。
 * 时间复杂度为O(N)。
 * 随机选取m个数，每个数被选中的概率为m/N。
 *
 * 算法思路大致如下：
 *
 * 1. 如果接收的数据量小于m，则依次放入蓄水池。
 * 2. 当接收到第i个数据时，i >= m，在[0, i]范围内取以随机数d，若d的落在[0, m-1]范围内，则用接收到的第i个数据替换蓄水池中的第d个数据。
 * 3. 重复步骤2。
 *
 * 证明：
 *
 * 对于任意的 i，第 i 个元素进入蓄水池的概率为 m / i
 * 而在蓄水池内每个元素被替换的概率为 1/m;
 * 因此在第 i 轮第j个元素被替换的概率为 (m/i) * (1/m) = 1/i。

 * 接下来用数学归纳法来证明，当循环结束时每个元素进入蓄水池的概率为 m / N.
 *
 * 假设在 (i-1) 次迭代后，任意一个元素进入 蓄水池的概率为 m / (i-1)。
 * 有上面的结论，在第 i 次迭代时，该元素被替换的概率为 1 / i， 那么其不被替换的概率则为 1 - 1/i = (i-1)/i；
 * 在第i 次迭代后，该元素在蓄水池内的概率为 m / (i-1) * (i-1)/i = m / i. 归纳部分结束。
 *
 * 因此当循环结束时，每个元素进入蓄水池的概率为 m / N. 命题得证。
 *
 * https://www.jianshu.com/p/7a9ea6ece2af
 */


public class ReservoirSampling {

    /**
     *
     * @param dataStream 数据流
     * @param m 采样数
     * @return
     */
    public int[] sample(int[] dataStream, int m) {
        int[] res = new int[m];
        Random rand = new Random();

        for (int i = 0; i < m; i++) {
            res[i] = dataStream[i];
        }

        for (int i = m; i < dataStream.length; i++) {
            // 随机获得一个[0, i]内的随机整数
            int d = rand.nextInt(i + 1);
            // 如果随机整数落在[0, m-1]范围内，则替换蓄水池中的元素
            if (d < m) {
                res[d] = dataStream[i];
            }
        }

        return res;
    }

}
